import torch
from torch import nn
from torchvision import models


def getResNet18():
    net = models.get_model("resnet18")
    # net = models.resnet18()
    # 修改第一层卷积层和最后的全连接层，使其适应CIFAR-10的分类任务
    net.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
    # 修改分类输出
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 10)
    return net

def getResNet50():
    net = models.get_model("resnet50")
    # net = models.resnet18()
    # 修改第一层卷积层和最后的全连接层，使其适应CIFAR-10的分类任务
    net.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
    # 修改分类输出
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 10)
    return net

def get_resnet101():
   # 加载预训练的ResNet101模型
   net = models.resnet101(pretrained=True)
   # 修改第一层卷积层，使其适应32x32尺寸的输入
   net.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
   # 修改分类输出
   num_ftrs = net.fc.in_features
   net.fc = nn.Linear(num_ftrs, 10)  # CIFAR-10有10个类别
   return net

def get_resnet152():
    # 加载预训练的ResNet152模型
    net = models.resnet152(pretrained=True)

    # 修改第一层卷积层，使其适应32x32尺寸的输入
    net.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)

    # 修改分类输出
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 10)  # CIFAR-10有10个类别

    return net
def getDensenet201():
    # 加载预训练的DenseNet-201模型
    model = models.densenet201(pretrained=True)
    # 修改分类输出
    num_features = model.classifier.in_features
    model.classifier = torch.nn.Linear(num_features, 10)  # num_classes是你任务的类别数
    return model

def getAvailableDevice():
    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

